import org.apache.spark.SparkContext;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.util.LongAccumulator;
import redis.clients.jedis.Jedis;
import utils.RedisInstance;

import java.util.concurrent.atomic.AtomicLong;

public class WriteRedis {
    private static final int totalQps=10;

    public static void main(String[] args) {

        SparkSession spark = SparkSession.builder()
                .appName("RDDToDataset")
                .master("local[*]")
                .getOrCreate();
        //todo
        /**
         * 1.source
         *
         * 2.tranform
         *
         * 3.sink
         */
        JavaRDD<String> lines = spark.read().textFile("src/main/resources/input").javaRDD();

        JavaRDD<Row> result = lines.map(new Function<String, Row>() {
            @Override
            public Row call(String line) throws Exception {
                String[] str = line.split(" ");
                return RowFactory.create(str[0],str[1]);
            }
        });
        int partition = result.getNumPartitions();
        //System.out.println("par1:" + result.getNumPartitions());
        //result=result.repartition(2);
        //System.out.println("par2:" + result.getNumPartitions());
        int requiredQps = totalQps/partition;

        System.out.println("total:"+result.count());
        result.foreachPartition(it -> {
            Jedis jedis = RedisInstance.getInstance();
            System.out.println(it.hashCode());
            AtomicLong atomicLong = new AtomicLong();
            long start = System.currentTimeMillis();
            it.forEachRemaining(v -> {
                //System.out.println(v.getString(0)+":"+v.getString(1));
                atomicLong.incrementAndGet();
                qpsControll(start,requiredQps,atomicLong,it.hashCode());
                jedis.set(v.getString(0),v.getString(1));
                    }
            );
        });
        //System.out.println(count.value());
        spark.stop();
    }

    private static void qpsControll(long start, int requiredQps, AtomicLong count,int x) {
        //System.out.println("current count:"+x+":"+ count.get());
        long actualQps = 1000*count.get()/(System.currentTimeMillis()-start);
        System.out.println(x+":"+actualQps);
        if (actualQps>(long)requiredQps){
            System.out.println("=====stop=====");
            try {
                Thread.sleep(3000);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }


}
